{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jstjohn/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "beViQ8vrM6LY",
        "outputId": "1a821c07-6701-457d-ddd3-a3f52a03d4a6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'enformer-pytorch'...\n",
            "remote: Enumerating objects: 643, done.\u001b[K\n",
            "remote: Counting objects: 100% (132/132), done.\u001b[K\n",
            "remote: Compressing objects: 100% (117/117), done.\u001b[K\n",
            "remote: Total 643 (delta 28), reused 28 (delta 13), pack-reused 511\u001b[K\n",
            "Receiving objects: 100% (643/643), 8.88 MiB | 3.09 MiB/s, done.\n",
            "Resolving deltas: 100% (439/439), done.\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/lucidrains/enformer-pytorch.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PMXWgyEfNT_s",
        "outputId": "804379ef-37af-481c-c870-b3e9c7b8e2bb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Processing /content/enformer-pytorch\n",
            "\u001b[33m  DEPRECATION: A future pip version will change local packages to be built in-place without first copying to a temporary directory. We recommend you use --use-feature=in-tree-build to test your packages with this new behavior before it becomes the default.\n",
            "   pip 21.3 will remove support for this functionality. You can find discussion regarding this at https://github.com/pypa/pip/issues/7555.\u001b[0m\n",
            "Collecting einops>=0.3\n",
            "  Downloading einops-0.4.1-py3-none-any.whl (28 kB)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from enformer-pytorch==0.5.1) (1.21.6)\n",
            "Requirement already satisfied: torch>=1.6 in /usr/local/lib/python3.7/dist-packages (from enformer-pytorch==0.5.1) (1.11.0+cu113)\n",
            "Collecting polars\n",
            "  Downloading polars-0.13.40-cp37-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (11.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 11.9 MB 19.1 MB/s \n",
            "\u001b[?25hCollecting pyfaidx\n",
            "  Downloading pyfaidx-0.7.0.tar.gz (102 kB)\n",
            "\u001b[K     |████████████████████████████████| 102 kB 70.3 MB/s \n",
            "\u001b[?25hRequirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from enformer-pytorch==0.5.1) (3.13)\n",
            "Collecting transformers\n",
            "  Downloading transformers-4.19.2-py3-none-any.whl (4.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.2 MB 56.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.6->enformer-pytorch==0.5.1) (4.2.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from pyfaidx->enformer-pytorch==0.5.1) (1.15.0)\n",
            "Requirement already satisfied: setuptools>=0.7 in /usr/local/lib/python3.7/dist-packages (from pyfaidx->enformer-pytorch==0.5.1) (57.4.0)\n",
            "Collecting tokenizers!=0.11.3,<0.13,>=0.11.1\n",
            "  Downloading tokenizers-0.12.1-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (6.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 6.6 MB 44.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.7/dist-packages (from transformers->enformer-pytorch==0.5.1) (4.64.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers->enformer-pytorch==0.5.1) (3.7.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers->enformer-pytorch==0.5.1) (2019.12.20)\n",
            "Collecting huggingface-hub<1.0,>=0.1.0\n",
            "  Downloading huggingface_hub-0.7.0-py3-none-any.whl (86 kB)\n",
            "\u001b[K     |████████████████████████████████| 86 kB 6.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from transformers->enformer-pytorch==0.5.1) (4.11.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.7/dist-packages (from transformers->enformer-pytorch==0.5.1) (21.3)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers->enformer-pytorch==0.5.1) (2.23.0)\n",
            "Collecting pyyaml\n",
            "  Downloading PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (596 kB)\n",
            "\u001b[K     |████████████████████████████████| 596 kB 64.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=20.0->transformers->enformer-pytorch==0.5.1) (3.0.9)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->transformers->enformer-pytorch==0.5.1) (3.8.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers->enformer-pytorch==0.5.1) (2022.5.18.1)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers->enformer-pytorch==0.5.1) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers->enformer-pytorch==0.5.1) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers->enformer-pytorch==0.5.1) (2.10)\n",
            "Building wheels for collected packages: enformer-pytorch, pyfaidx\n",
            "  Building wheel for enformer-pytorch (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for enformer-pytorch: filename=enformer_pytorch-0.5.1-py3-none-any.whl size=11587 sha256=3888ddb9602b37eedeb4b73f60c5d56609dd500e54536df6fb1bb50f481fe6d5\n",
            "  Stored in directory: /root/.cache/pip/wheels/e9/30/2d/17bcf281153d214e2a57d1ed68544840825283066370c73eda\n",
            "  Building wheel for pyfaidx (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyfaidx: filename=pyfaidx-0.7.0-py3-none-any.whl size=27697 sha256=464766257aa70062f3ac14de5005e39af8e5375eb2be08e6f32e0a05dd994a12\n",
            "  Stored in directory: /root/.cache/pip/wheels/df/6b/ce/46374a70af569061fa10a6c16525b0d8efe2d9a4069f8a144a\n",
            "Successfully built enformer-pytorch pyfaidx\n",
            "Installing collected packages: pyyaml, tokenizers, huggingface-hub, transformers, pyfaidx, polars, einops, enformer-pytorch\n",
            "  Attempting uninstall: pyyaml\n",
            "    Found existing installation: PyYAML 3.13\n",
            "    Uninstalling PyYAML-3.13:\n",
            "      Successfully uninstalled PyYAML-3.13\n",
            "Successfully installed einops-0.4.1 enformer-pytorch-0.5.1 huggingface-hub-0.7.0 polars-0.13.40 pyfaidx-0.7.0 pyyaml-6.0 tokenizers-0.12.1 transformers-4.19.2\n"
          ]
        }
      ],
      "source": [
        "!cd enformer-pytorch && pip install ."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "0XbZipT7O0yD"
      },
      "outputs": [],
      "source": [
        "!pip install torchmetrics kipoiseq==0.5.2 BioPython --quiet > /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H8miDmi7OMK1",
        "outputId": "4249c35d-2c6f-4c0f-cf70-16d74da61f24"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Copying gs://basenji_barnyard/hg38.ml.fa.gz...\n",
            "/ [1/1 files][839.8 MiB/839.8 MiB] 100% Done  58.5 MiB/s ETA 00:00:00           \n",
            "Operation completed over 1 objects/839.8 MiB.                                    \n",
            "Copying gs://basenji_barnyard/mm10.ml.fa.gz...\n",
            "/ [1/1 files][800.8 MiB/800.8 MiB] 100% Done  65.9 MiB/s ETA 00:00:00           \n",
            "Operation completed over 1 objects/800.8 MiB.                                    \n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import os \n",
        "import json\n",
        "import pandas as pd\n",
        "import pyfaidx\n",
        "import kipoiseq\n",
        "import functools\n",
        "from kipoiseq import Interval\n",
        "\n",
        "SEQUENCE_LENGTH = 196_608\n",
        "BIN_SIZE = 128\n",
        "TARGET_LENGTH = 896\n",
        "import os\n",
        "fasta_dir = \"/root/data/\"\n",
        "!mkdir -p {fasta_dir}\n",
        "human_fasta_f = 'hg38.ml.fa.gz'\n",
        "mouse_fasta_f = 'mm10.ml.fa.gz'\n",
        "human_fasta_gz_path = f\"{fasta_dir}/{human_fasta_f}\"\n",
        "mouse_fasta_gz_path = f\"{fasta_dir}/{mouse_fasta_f}\"\n",
        "human_fasta_path = human_fasta_gz_path.rstrip(\".gz\")\n",
        "mouse_fasta_path = mouse_fasta_gz_path.rstrip(\".gz\")\n",
        "\n",
        "if not os.path.isfile(human_fasta_path):\n",
        "  !gsutil -m cp -n gs://basenji_barnyard/{human_fasta_f} {human_fasta_gz_path}\n",
        "  !gunzip {human_fasta_gz_path}\n",
        "if not os.path.isfile(mouse_fasta_path):\n",
        "  !gsutil -m cp -n gs://basenji_barnyard/{mouse_fasta_f} {mouse_fasta_gz_path}\n",
        "  !gunzip {mouse_fasta_gz_path}\n",
        "\n",
        "class FastaStringExtractor:\n",
        "    \n",
        "    def __init__(self, fasta_file):\n",
        "        self.fasta = pyfaidx.Fasta(fasta_file)\n",
        "        self._chromosome_sizes = {k: len(v) for k, v in self.fasta.items()}\n",
        "\n",
        "    def extract(self, interval: Interval, **kwargs) -> str:\n",
        "        # Truncate interval if it extends beyond the chromosome lengths.\n",
        "        chromosome_length = self._chromosome_sizes[interval.chrom]\n",
        "        trimmed_interval = Interval(interval.chrom,\n",
        "                                    max(interval.start, 0),\n",
        "                                    min(interval.end, chromosome_length),\n",
        "                                    )\n",
        "        # pyfaidx wants a 1-based interval\n",
        "        sequence = str(self.fasta.get_seq(trimmed_interval.chrom,\n",
        "                                          trimmed_interval.start + 1,\n",
        "                                          trimmed_interval.stop).seq).upper()\n",
        "        # Fill truncated values with N's.\n",
        "        pad_upstream = 'N' * max(-interval.start, 0)\n",
        "        pad_downstream = 'N' * max(interval.end - chromosome_length, 0)\n",
        "        return pad_upstream + sequence + pad_downstream\n",
        "\n",
        "    def close(self):\n",
        "        return self.fasta.close()\n",
        "\n",
        "\n",
        "class BasenjiDataSet(torch.utils.data.IterableDataset):\n",
        "  @staticmethod\n",
        "  def get_organism_path(organism):\n",
        "    return os.path.join('gs://basenji_barnyard/data', organism)\n",
        "  @classmethod\n",
        "  def get_metadata(cls, organism):\n",
        "    # Keys:\n",
        "    # num_targets, train_seqs, valid_seqs, test_seqs, seq_length,\n",
        "    # pool_width, crop_bp, target_length\n",
        "    path = os.path.join(cls.get_organism_path(organism), 'statistics.json')\n",
        "    with tf.io.gfile.GFile(path, 'r') as f:\n",
        "      return json.load(f)\n",
        "  @staticmethod\n",
        "  def one_hot_encode(sequence):\n",
        "    return kipoiseq.transforms.functional.one_hot_dna(sequence).astype(np.float32)\n",
        "\n",
        "  @classmethod\n",
        "  def get_tfrecord_files(cls, organism, subset):\n",
        "    # Sort the values by int(*).\n",
        "    return sorted(tf.io.gfile.glob(os.path.join(\n",
        "        cls.get_organism_path(organism), 'tfrecords', f'{subset}-*.tfr'\n",
        "      )), key=lambda x: int(x.split('-')[-1].split('.')[0]))\n",
        "  \n",
        "  @property\n",
        "  def num_channels(self):\n",
        "    metadata = self.get_metadata(self.organism)\n",
        "    return metadata['num_targets']\n",
        "\n",
        "  @staticmethod\n",
        "  def deserialize(serialized_example, metadata):\n",
        "    \"\"\"Deserialize bytes stored in TFRecordFile.\"\"\"\n",
        "    # Deserialization\n",
        "    feature_map = {\n",
        "        'sequence': tf.io.FixedLenFeature([], tf.string),  # Ignore this, resize our own bigger one\n",
        "        'target': tf.io.FixedLenFeature([], tf.string),\n",
        "    }\n",
        "    example = tf.io.parse_example(serialized_example, feature_map)\n",
        "    sequence = tf.io.decode_raw(example['sequence'], tf.bool)\n",
        "    sequence = tf.reshape(sequence, (metadata['seq_length'], 4))\n",
        "    sequence = tf.cast(sequence, tf.float32)\n",
        "\n",
        "    target = tf.io.decode_raw(example['target'], tf.float16)\n",
        "    target = tf.reshape(target,\n",
        "                        (metadata['target_length'], metadata['num_targets']))\n",
        "    target = tf.cast(target, tf.float32)\n",
        "\n",
        "    return {'sequence_old': sequence,\n",
        "            'target': target}\n",
        "\n",
        "  @classmethod\n",
        "  def get_dataset(cls, organism, subset, num_threads=8):\n",
        "    metadata = cls.get_metadata(organism)\n",
        "    dataset = tf.data.TFRecordDataset(cls.get_tfrecord_files(organism, subset),\n",
        "                                      compression_type='ZLIB',\n",
        "                                      num_parallel_reads=num_threads).map(\n",
        "                                          functools.partial(cls.deserialize, metadata=metadata)\n",
        "                                      )\n",
        "    return dataset\n",
        "\n",
        "  def __init__(self, organism:str, subset:str, seq_len:int, fasta_path:str, n_to_test:int = -1):\n",
        "    assert subset in {\"train\", \"valid\", \"test\"}\n",
        "    assert organism in {\"human\", \"mouse\"}\n",
        "    self.organism = organism\n",
        "    self.subset = subset\n",
        "    self.base_dir = self.get_organism_path(organism)\n",
        "    self.seq_len = seq_len\n",
        "    self.fasta_reader = FastaStringExtractor(fasta_path)\n",
        "    self.n_to_test = n_to_test\n",
        "    with tf.io.gfile.GFile(f\"{self.base_dir}/sequences.bed\", 'r') as f:\n",
        "      region_df = pd.read_csv(f, sep=\"\\t\", header=None)\n",
        "      region_df.columns = ['chrom', 'start', 'end', 'subset']\n",
        "      self.region_df = region_df.query('subset==@subset').reset_index(drop=True)\n",
        "      \n",
        "  def __iter__(self):\n",
        "    worker_info = torch.utils.data.get_worker_info()\n",
        "    assert worker_info is None, \"Only support single process loading\"\n",
        "    # If num_threads > 1, the following will actually shuffle the inputs! luckily we catch this with the sequence comparison\n",
        "    basenji_iterator = self.get_dataset(self.organism, self.subset, num_threads=1).as_numpy_iterator()\n",
        "    for i, records in enumerate(basenji_iterator):\n",
        "      loc_row = self.region_df.iloc[i]\n",
        "      target_interval = Interval(loc_row['chrom'], loc_row['start'], loc_row['end'])\n",
        "      sequence_one_hot = self.one_hot_encode(self.fasta_reader.extract(target_interval.resize(self.seq_len)))\n",
        "      if self.n_to_test >= 0 and i < self.n_to_test:\n",
        "        old_sequence_onehot = records[\"sequence_old\"]\n",
        "        if old_sequence_onehot.shape[0] > sequence_one_hot.shape[0]:\n",
        "          diff = old_sequence_onehot.shape[0] - sequence_one_hot.shape[0]\n",
        "          trim = diff//2\n",
        "          np.testing.assert_equal(old_sequence_onehot[trim:(-trim)], sequence_one_hot)\n",
        "        elif sequence_one_hot.shape[0] > old_sequence_onehot.shape[0]:\n",
        "          diff = sequence_one_hot.shape[0] - old_sequence_onehot.shape[0]\n",
        "          trim = diff//2\n",
        "          np.testing.assert_equal(old_sequence_onehot, sequence_one_hot[trim:(-trim)])\n",
        "        else:\n",
        "          np.testing.assert_equal(old_sequence_onehot, sequence_one_hot)\n",
        "      yield {\n",
        "          \"sequence\": sequence_one_hot,\n",
        "          \"target\": records[\"target\"],\n",
        "      }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 81,
          "referenced_widgets": [
            "a76d636b937c4b5691d36cb6b7ec7b17",
            "e06f0950617847e7a894afc8cc8ec69f",
            "fd6815290f9e45f69d9ac29826de71ef",
            "109781990b25433daa411cafa44c662a",
            "69789023fe8544ff8a0d7540de793041",
            "ef2d9fda87da4dbc825f4776a982856d",
            "20f0fc109b5249a3bd0e5738f6b8c06f",
            "7ab66e50e29a48b38a1d5c1820422a6a",
            "b0b4b1ebc3ac4aaf91b0112969fc64f9",
            "38e1b0815aaf4c01b1988f73fc3c38d2",
            "c177e2c35da948918a29fdcea6840d89",
            "beeeb50adabb48978b1c502d799713be",
            "a34e3298896c454abb1b0dda7c225c2b",
            "579ec7740cbc455a9c6d3e1bc4442612",
            "e486e3955a574b818227964293fa76e1",
            "44a8b215df314ba9beb66b44d2e13d29",
            "36a42d5cd33541b5aa24db50b03cef9a",
            "e47a62c2ab1a4fcea674274f3435f031",
            "02df913b81f3473cbbe0189880e20586",
            "9d3909432e2d42a5849427f6d1ad1371",
            "ec0bf3f9df7c40b8b64aef5ac62f0728",
            "7ba9a6bb917746d2bdb7143843480038"
          ]
        },
        "id": "W_QtUyQXNfgJ",
        "outputId": "2b58967c-2b79-439f-a901-687a164282d3"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading:   0%|          | 0.00/464 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a76d636b937c4b5691d36cb6b7ec7b17"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading:   0%|          | 0.00/959M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "beeeb50adabb48978b1c502d799713be"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "import torch\n",
        "from enformer_pytorch import Enformer\n",
        "\n",
        "model = Enformer.from_pretrained(\"EleutherAI/enformer-official-rough\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "3wgrVVQIRB-w"
      },
      "outputs": [],
      "source": [
        "model = model.eval().cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "2MCoTxFDVPuV"
      },
      "outputs": [],
      "source": [
        "from enformer_pytorch.metrics import MeanPearsonCorrCoefPerChannel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0GFJBmbLQ1-R",
        "outputId": "d10a6a75-d141-4d45-bf76-7f5e2d23c6ec"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [01:22<00:00,  1.21it/s]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(0.6270)"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ],
      "source": [
        "from tqdm import tqdm\n",
        "from torchmetrics.regression.pearson import PearsonCorrCoef\n",
        "def compute_correlation(model, organism:str=\"human\", subset:str=\"valid\", max_steps=-1):\n",
        "  fasta_path = human_fasta_path if organism == \"human\" else mouse_fasta_path\n",
        "  ds = BasenjiDataSet(organism, subset, SEQUENCE_LENGTH, fasta_path)\n",
        "  total = len(ds.region_df) # number of records\n",
        "  dl = torch.utils.data.DataLoader(ds, num_workers=0, batch_size=1)\n",
        "  corr_coef = MeanPearsonCorrCoefPerChannel(n_channels=ds.num_channels)\n",
        "  n_steps = total if max_steps <= 0 else max_steps\n",
        "  for i,batch in enumerate(tqdm(dl, total=n_steps)):\n",
        "    if max_steps > 0 and i >= max_steps:\n",
        "      break\n",
        "    batch_gpu = {k:v.to(model.device) for k,v in batch.items()}\n",
        "    sequence = batch_gpu['sequence']\n",
        "    target = batch_gpu['target']\n",
        "    with torch.no_grad():\n",
        "      pred = model(sequence)[organism]\n",
        "      corr_coef(preds=pred.cpu(), target=target.cpu())\n",
        "  return corr_coef.compute().mean()\n",
        "compute_correlation(model, organism=\"human\", subset=\"valid\", max_steps=100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true,
          "base_uri": "https://localhost:8080/"
        },
        "id": "raZi83xyg8SV",
        "outputId": "c00dae32-0838-4b4d-a9a1-be6040133011"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torchmetrics/utilities/prints.py:36: UserWarning: Torchmetrics v0.9 introduced a new argument class property called `full_state_update` that has\n",
            "                not been set for this class (MeanPearsonCorrCoefPerChannel). The property determines if `update` by\n",
            "                default needs access to the full metric state. If this is not the case, significant speedups can be\n",
            "                achieved and we recommend setting this to `False`.\n",
            "                We provide an checking function\n",
            "                `from torchmetrics.utilities import check_forward_no_full_state`\n",
            "                that can be used to check if the `full_state_update=True` (old and potential slower behaviour,\n",
            "                default for now) or if `full_state_update=False` can be used safely.\n",
            "                \n",
            "  warnings.warn(*args, **kwargs)\n",
            "100%|██████████| 2213/2213 [30:51<00:00,  1.20it/s]\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "tensor(0.6252)"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "compute_correlation(model, organism=\"human\", subset=\"valid\", max_steps=-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "PtWriiqConiR",
        "outputId": "20acb7f9-3ccc-425c-83d6-99ad88100cba"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torchmetrics/utilities/prints.py:36: UserWarning: Torchmetrics v0.9 introduced a new argument class property called `full_state_update` that has\n",
            "                not been set for this class (MeanPearsonCorrCoefPerChannel). The property determines if `update` by\n",
            "                default needs access to the full metric state. If this is not the case, significant speedups can be\n",
            "                achieved and we recommend setting this to `False`.\n",
            "                We provide an checking function\n",
            "                `from torchmetrics.utilities import check_forward_no_full_state`\n",
            "                that can be used to check if the `full_state_update=True` (old and potential slower behaviour,\n",
            "                default for now) or if `full_state_update=False` can be used safely.\n",
            "                \n",
            "  warnings.warn(*args, **kwargs)\n",
            "100%|██████████| 1937/1937 [20:05<00:00,  1.61it/s]\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "tensor(0.6503)"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "compute_correlation(model, organism=\"human\", subset=\"test\", max_steps=-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "oRMRZUMHo5XX",
        "outputId": "68b054a6-e6b7-4949-9866-ad4b9615fc9e"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torchmetrics/utilities/prints.py:36: UserWarning: Torchmetrics v0.9 introduced a new argument class property called `full_state_update` that has\n",
            "                not been set for this class (MeanPearsonCorrCoefPerChannel). The property determines if `update` by\n",
            "                default needs access to the full metric state. If this is not the case, significant speedups can be\n",
            "                achieved and we recommend setting this to `False`.\n",
            "                We provide an checking function\n",
            "                `from torchmetrics.utilities import check_forward_no_full_state`\n",
            "                that can be used to check if the `full_state_update=True` (old and potential slower behaviour,\n",
            "                default for now) or if `full_state_update=False` can be used safely.\n",
            "                \n",
            "  warnings.warn(*args, **kwargs)\n",
            "100%|██████████| 34021/34021 [4:58:22<00:00,  1.90it/s]\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "tensor(0.7415)"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "compute_correlation(model, organism=\"human\", subset=\"train\", max_steps=-1)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "background_execution": "on",
      "collapsed_sections": [],
      "name": "evaluate_enformer_pytorch_correlation.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyP5rds3GvFsVbB6KySflndx",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "a76d636b937c4b5691d36cb6b7ec7b17": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e06f0950617847e7a894afc8cc8ec69f",
              "IPY_MODEL_fd6815290f9e45f69d9ac29826de71ef",
              "IPY_MODEL_109781990b25433daa411cafa44c662a"
            ],
            "layout": "IPY_MODEL_69789023fe8544ff8a0d7540de793041"
          }
        },
        "e06f0950617847e7a894afc8cc8ec69f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ef2d9fda87da4dbc825f4776a982856d",
            "placeholder": "​",
            "style": "IPY_MODEL_20f0fc109b5249a3bd0e5738f6b8c06f",
            "value": "Downloading: 100%"
          }
        },
        "fd6815290f9e45f69d9ac29826de71ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7ab66e50e29a48b38a1d5c1820422a6a",
            "max": 464,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_b0b4b1ebc3ac4aaf91b0112969fc64f9",
            "value": 464
          }
        },
        "109781990b25433daa411cafa44c662a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_38e1b0815aaf4c01b1988f73fc3c38d2",
            "placeholder": "​",
            "style": "IPY_MODEL_c177e2c35da948918a29fdcea6840d89",
            "value": " 464/464 [00:00&lt;00:00, 12.3kB/s]"
          }
        },
        "69789023fe8544ff8a0d7540de793041": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ef2d9fda87da4dbc825f4776a982856d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "20f0fc109b5249a3bd0e5738f6b8c06f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7ab66e50e29a48b38a1d5c1820422a6a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b0b4b1ebc3ac4aaf91b0112969fc64f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "38e1b0815aaf4c01b1988f73fc3c38d2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c177e2c35da948918a29fdcea6840d89": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "beeeb50adabb48978b1c502d799713be": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a34e3298896c454abb1b0dda7c225c2b",
              "IPY_MODEL_579ec7740cbc455a9c6d3e1bc4442612",
              "IPY_MODEL_e486e3955a574b818227964293fa76e1"
            ],
            "layout": "IPY_MODEL_44a8b215df314ba9beb66b44d2e13d29"
          }
        },
        "a34e3298896c454abb1b0dda7c225c2b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_36a42d5cd33541b5aa24db50b03cef9a",
            "placeholder": "​",
            "style": "IPY_MODEL_e47a62c2ab1a4fcea674274f3435f031",
            "value": "Downloading: 100%"
          }
        },
        "579ec7740cbc455a9c6d3e1bc4442612": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_02df913b81f3473cbbe0189880e20586",
            "max": 1005149571,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_9d3909432e2d42a5849427f6d1ad1371",
            "value": 1005149571
          }
        },
        "e486e3955a574b818227964293fa76e1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ec0bf3f9df7c40b8b64aef5ac62f0728",
            "placeholder": "​",
            "style": "IPY_MODEL_7ba9a6bb917746d2bdb7143843480038",
            "value": " 959M/959M [00:36&lt;00:00, 58.7MB/s]"
          }
        },
        "44a8b215df314ba9beb66b44d2e13d29": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "36a42d5cd33541b5aa24db50b03cef9a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e47a62c2ab1a4fcea674274f3435f031": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "02df913b81f3473cbbe0189880e20586": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9d3909432e2d42a5849427f6d1ad1371": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "ec0bf3f9df7c40b8b64aef5ac62f0728": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7ba9a6bb917746d2bdb7143843480038": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}